#!/bin/bash

# 通用多模态信息抽取 - 增强训练监控脚本
# 提供实时训练状态显示和日志监控

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 项目路径
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_DIR="${PROJECT_ROOT}/logs/enhanced_training"

# 创建日志目录
mkdir -p "${LOG_DIR}"

# 打印带颜色的消息
print_message() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
}

# 打印标题
print_title() {
    echo -e "${PURPLE}"
    echo "======================================================================"
    echo "             通用多模态信息抽取 - 增强训练监控系统"
    echo "======================================================================"
    echo -e "${NC}"
}

# 检查环境
check_environment() {
    print_message $BLUE "🔍 检查训练环境..."
    
    # 检查Python环境
    if ! command -v python &> /dev/null; then
        print_message $RED "❌ Python未安装"
        exit 1
    fi
    
    # 检查Swift框架
    if ! python -c "import swift" 2>/dev/null; then
        print_message $RED "❌ Swift框架未安装，请运行: pip install ms-swift"
        exit 1
    fi
    
    # 检查CUDA
    if ! python -c "import torch; print('CUDA可用' if torch.cuda.is_available() else 'CUDA不可用')" 2>/dev/null | grep -q "CUDA可用"; then
        print_message $RED "❌ CUDA不可用"
        exit 1
    fi
    
    # 检查训练数据
    if [ ! -f "${PROJECT_ROOT}/data/train_data_mixed_train.jsonl" ]; then
        print_message $YELLOW "⚠️  混合训练数据不存在，使用原始训练数据"
        if [ ! -f "${PROJECT_ROOT}/data/train_data_train.jsonl" ]; then
            print_message $RED "❌ 未找到训练数据文件"
            print_message $YELLOW "请先运行: python scripts/generate_sample_data.py"
            exit 1
        fi
    fi
    
    print_message $GREEN "✅ 环境检查完成"
}

# 显示GPU信息
show_gpu_info() {
    print_message $BLUE "🖥️  GPU信息:"
    python -c "
import torch
if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        gpu = torch.cuda.get_device_properties(i)
        memory_gb = gpu.total_memory / 1024**3
        print(f'  GPU {i}: {gpu.name} ({memory_gb:.1f}GB)')
        
    # 显示当前GPU使用情况
    import subprocess
    try:
        result = subprocess.run(['nvidia-smi', '--query-gpu=index,memory.used,memory.total', '--format=csv,noheader,nounits'], 
                              capture_output=True, text=True)
        if result.returncode == 0:
            print('  当前GPU内存使用:')
            for line in result.stdout.strip().split('\n'):
                if line.strip():
                    parts = line.split(', ')
                    if len(parts) >= 3:
                        gpu_id, used, total = parts[0], parts[1], parts[2]
                        usage_percent = (int(used) / int(total)) * 100
                        print(f'    GPU {gpu_id}: {used}MB / {total}MB ({usage_percent:.1f}%)')
    except:
        pass
else:
    print('  ❌ 未检测到可用GPU')
"
}

# 智能训练监控函数 - 支持自动检测训练完成
monitor_training() {
    local log_file=$1
    local stage_name=$2
    local train_pid=$3  # 训练进程PID参数
    
    print_message $CYAN "📊 开始监控${stage_name}训练进度..."
    print_message $YELLOW "日志文件: ${log_file}"
    print_message $YELLOW "训练进程PID: ${train_pid}"
    
    # 等待日志文件创建
    local wait_count=0
    while [ ! -f "${log_file}" ] && [ $wait_count -lt 30 ]; do
        sleep 2
        ((wait_count++))
        print_message $YELLOW "等待训练开始... (${wait_count}/30)"
        
        # 检查进程是否还在运行
        if [ -n "$train_pid" ] && ! kill -0 "$train_pid" 2>/dev/null; then
            print_message $RED "❌ 训练进程已停止，但未找到日志文件"
            return 1
        fi
    done
    
    if [ ! -f "${log_file}" ]; then
        print_message $RED "❌ 日志文件未创建，可能训练启动失败"
        return 1
    fi
    
    # 智能监控变量
    local last_size=0
    local no_change_count=0
    local max_no_change=15  # 降低等待时间
    local training_completed=false
    local completion_wait_count=0
    local max_completion_wait=30  # 训练完成后最多等待30次
    
    print_message $BLUE "📋 开始实时监控训练进度..."
    
    while true; do
        # 检查训练进程是否还在运行
        if [ -n "$train_pid" ] && ! kill -0 "$train_pid" 2>/dev/null; then
            print_message $GREEN "✅ 训练进程已完成，停止监控"
            break
        fi
        
        # 检查日志文件大小变化
        if [ -f "${log_file}" ]; then
            local current_size=$(stat -f%z "${log_file}" 2>/dev/null || stat -c%s "${log_file}" 2>/dev/null || echo "0")
            
            if [ "$current_size" -gt "$last_size" ]; then
                # 有新内容，显示新增的日志行
                local new_content=$(tail -c +$((last_size + 1)) "${log_file}")
                
                # 显示新的日志内容
                echo "$new_content" | while read line; do
                    if [ -n "$line" ]; then
                        # 高亮显示重要信息
                        if echo "$line" | grep -iE "(loss|epoch|step|lr|gpu|error|warning|完成|成功|失败|Training completed|训练完成|100%|train_runtime|global_step)" >/dev/null; then
                            if echo "$line" | grep -iE "(error|失败|Error|Exception)" >/dev/null; then
                                print_message $RED "🚨 $line"
                            elif echo "$line" | grep -iE "(完成|成功|completed|finished|success|100%)" >/dev/null; then
                                print_message $GREEN "✅ $line"
                            elif echo "$line" | grep -iE "(loss|epoch|step|/s]|token|accuracy|train_runtime)" >/dev/null; then
                                print_message $CYAN "📈 $line"
                            else
                                print_message $YELLOW "ℹ️  $line"
                            fi
                        else
                            echo "$line"
                        fi
                    fi
                done
                
                # 检查是否出现训练完成的标志
                if echo "$new_content" | grep -qE "(100%.*100%|train_runtime|End time of running main|训练任务成功完成)"; then
                    print_message $GREEN "✅ 检测到训练完成标志！"
                    training_completed=true
                fi
                
                last_size="$current_size"
                no_change_count=0
            else
                # 没有新内容
                ((no_change_count++))
                
                # 如果已检测到训练完成，等待进程退出
                if [ "$training_completed" = true ]; then
                    ((completion_wait_count++))
                    if [ $completion_wait_count -ge $max_completion_wait ]; then
                        print_message $YELLOW "⏰ 训练已完成但进程未退出，强制结束监控"
                        # 可以选择强制终止进程
                        if [ -n "$train_pid" ] && kill -0 "$train_pid" 2>/dev/null; then
                            print_message $YELLOW "⏳ 尝试温和终止训练进程..."
                            kill -TERM "$train_pid" 2>/dev/null
                            sleep 5
                            if kill -0 "$train_pid" 2>/dev/null; then
                                print_message $YELLOW "⚡ 强制终止训练进程..."
                                kill -KILL "$train_pid" 2>/dev/null
                            fi
                        fi
                        break
                    fi
                    print_message $BLUE "🗒️  训练已完成，等待进程退出... (${completion_wait_count}/${max_completion_wait})"
                elif [ $no_change_count -ge $max_no_change ]; then
                    print_message $YELLOW "⏰ 日志文件长时间无变化"
                    # 再次检查进程状态
                    if [ -n "$train_pid" ] && ! kill -0 "$train_pid" 2>/dev/null; then
                        print_message $GREEN "✅ 确认训练进程已完成"
                        break
                    fi
                    
                    # 检查日志文件中是否有完成标志
                    if [ -f "${log_file}" ]; then
                        local last_lines=$(tail -20 "${log_file}")
                        if echo "$last_lines" | grep -qE "(100%.*100%|train_runtime|End time of running main|训练任务成功完成)"; then
                            print_message $GREEN "✅ 日志显示训练已完成，结束监控"
                            training_completed=true
                            completion_wait_count=0  # 重置等待计数器
                        else
                            print_message $YELLOW "⏳ 训练进程仍在运行，继续等待..."
                            no_change_count=0  # 重置计数器
                        fi
                    fi
                fi
            fi
        fi
        
        sleep 2
    done
    
    # 显示最后的日志内容
    if [ -f "${log_file}" ]; then
        print_message $CYAN "📋 显示训练结束时的最后几行日志:"
        tail -10 "${log_file}" | while read line; do
            if [ -n "$line" ]; then
                print_message $YELLOW "   $line"
            fi
        done
    fi
    
    print_message $GREEN "🏁 ${stage_name}监控完成"
    return 0
}

# 启动TensorBoard监控
start_tensorboard() {
    local log_dir="${PROJECT_ROOT}/logs"
    
    if command -v tensorboard &> /dev/null; then
        print_message $BLUE "🔥 启动TensorBoard监控..."
        nohup tensorboard --logdir="${log_dir}" --port=6006 --host=0.0.0.0 > "${LOG_DIR}/tensorboard.log" 2>&1 &
        
        sleep 3
        if pgrep -f "tensorboard" > /dev/null; then
            print_message $GREEN "✅ TensorBoard已启动"
            print_message $CYAN "📊 访问地址: http://localhost:6006"
        else
            print_message $YELLOW "⚠️  TensorBoard启动失败，详见 ${LOG_DIR}/tensorboard.log"
        fi
    else
        print_message $YELLOW "⚠️  TensorBoard未安装，跳过可视化监控"
    fi
}

# 主训练函数
run_training() {
    local stage=$1
    local gpu_ids=$2
    
    # 时间戳
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    local training_log="${LOG_DIR}/training_${stage}_${timestamp}.log"
    
    print_message $BLUE "🚀 开始${stage}训练..."
    print_message $CYAN "使用GPU: ${gpu_ids}"
    
    # 启动训练（后台运行）
    cd "${PROJECT_ROOT}"
    nohup python scripts/universal_document_trainer.py --stage "${stage}" --gpu "${gpu_ids}" > "${training_log}" 2>&1 &
    local train_pid=$!
    
    print_message $GREEN "✅ 训练进程已启动 (PID: ${train_pid})"
    print_message $YELLOW "📝 主日志文件: ${training_log}"
    
    # 等待一下让训练初始化
    sleep 5
    
    # 查找Swift框架生成的详细日志
    local stage_log_dir="${PROJECT_ROOT}/logs/universal_document/stage${stage}"
    local stage_log_file="${stage_log_dir}/stage${stage}_training.log"
    
    # 监控训练进度
    if [ -f "${stage_log_file}" ]; then
        monitor_training "${stage_log_file}" "阶段${stage}"
    else
        print_message $YELLOW "⚠️  使用主日志文件监控"
        monitor_training "${training_log}" "阶段${stage}"
    fi
    
    # 等待训练完成
    wait $train_pid
    local exit_code=$?
    
    if [ $exit_code -eq 0 ]; then
        print_message $GREEN "🎉 阶段${stage}训练成功完成！"
        return 0
    else
        print_message $RED "❌ 阶段${stage}训练失败 (退出码: ${exit_code})"
        return 1
    fi
}

# 显示使用帮助
show_help() {
    echo "使用方法:"
    echo "  $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --stage <1|2|all>     指定训练阶段 (默认: interactive)"
    echo "  --gpu <gpu_ids>       指定GPU ID，用逗号分隔 (默认: interactive)"
    echo "  --monitor-only        仅启动监控，不执行训练"
    echo "  --tensorboard         启动TensorBoard"
    echo "  --help               显示此帮助信息"
    echo ""
    echo "交互模式示例:"
    echo "  $0                              # 进入交互式配置界面"
    echo ""
    echo "命令行模式示例:"
    echo "  $0 --stage 1 --gpu 0,1         # 仅运行阶段1，使用GPU 0和1"
    echo "  $0 --stage 2 --gpu 1           # 仅运行阶段2，使用GPU 1"
    echo "  $0 --tensorboard               # 启动TensorBoard监控"
}

# GPU配置选择界面
select_gpu_config() {
    print_message $BLUE "🖥️  GPU配置选择" >&2
    print_message $CYAN "当前GPU状态:" >&2
    
    # 显示GPU信息
    if command -v nvidia-smi &> /dev/null; then
        nvidia-smi --query-gpu=index,name,memory.total,memory.used --format=csv,noheader,nounits | head -3 | while read line; do
            if [ -n "$line" ]; then
                IFS=',' read -r gpu_id name total_mem used_mem <<< "$line"
                usage_percent=$(( used_mem * 100 / total_mem ))
                print_message $YELLOW "  GPU $gpu_id: $name (${used_mem}MB/${total_mem}MB, ${usage_percent}%)" >&2
            fi
        done
    else
        print_message $YELLOW "  ⚠️  nvidia-smi不可用，无法显示GPU状态" >&2
    fi
    
    echo ""
    print_message $CYAN "🔧 请选择您要使用的GPU配置方案:"
    echo ""
    print_message $GREEN "┌──────────────────────────────────────────────────────────────────┐"
    print_message $GREEN "│                        ⚙️  GPU配置选项说明                          │"
    print_message $GREEN "├──────────────────────────────────────────────────────────────────┤"
    print_message $YELLOW "│                                                                      │"
    print_message $YELLOW "│  选项1 📱 单GPU模式 (适合新手用户)                               │"
    print_message $YELLOW "│        - 使用显卡: GPU 0                                          │"
    print_message $YELLOW "│        - 显存需求: 约12-16GB                                      │"
    print_message $YELLOW "│        - 训练速度: 较慢，但稳定可靠                               │"
    print_message $YELLOW "│        - 适用场景: 第一次使用、硬件配置较低                       │"
    print_message $YELLOW "│                                                                      │"
    print_message $YELLOW "│  选项2 🔥 双GPU模式 (推荐大多数用户)                             │"
    print_message $YELLOW "│        - 使用显卡: GPU 0 + GPU 1                                  │"
    print_message $YELLOW "│        - 显存需求: 每张卡约8-12GB                                  │"
    print_message $YELLOW "│        - 训练速度: 中等，性价比高                                 │"
    print_message $YELLOW "│        - 适用场景: 一般用户，平衡性能与资源                       │"
    print_message $YELLOW "│                                                                      │"
    print_message $YELLOW "│  选项3 🚀 三GPU模式 (适合高级用户)                               │"
    print_message $YELLOW "│        - 使用显卡: GPU 0 + GPU 1 + GPU 2                         │"
    print_message $YELLOW "│        - 显存需求: 每张卡约6-10GB                                  │"
    print_message $YELLOW "│        - 训练速度: 快速，资源利用率高                             │"
    print_message $YELLOW "│        - 适用场景: 高级用户，追求训练速度                         │"
    print_message $YELLOW "│                                                                      │"
    print_message $YELLOW "│  选项4 ⚙️  自定义GPU配置 (专家选项)                               │"
    print_message $YELLOW "│        - 手动指定: 自由选择GPU编号组合                            │"
    print_message $YELLOW "│        - 灵活配置: 可指定任意GPU组合 (如: 1,3 或 0,2,3)            │"
    print_message $YELLOW "│        - 适用场景: 专家用户，特殊硬件配置                         │"
    print_message $YELLOW "│                                                                      │"
    print_message $YELLOW "│  选项5 📊 查看完整GPU信息 (信息查看)                              │"
    print_message $YELLOW "│        - 显示所有可用GPU的详细状态和使用情况                       │"
    print_message $YELLOW "│        - 帮助您了解当前硬件配置                                   │"
    print_message $YELLOW "│        - 选择后返回主菜单继续配置                                 │"
    print_message $YELLOW "│                                                                      │"
    print_message $GREEN "└──────────────────────────────────────────────────────────────────┘"
    echo ""
    print_message $CYAN "💡 提示: 建议根据您的显卡显存大小选择合适的配置"
    print_message $CYAN "   - 显存≥16GB: 推荐选择1 (单GPU)  "
    print_message $CYAN "   - 显存8-16GB: 推荐选择2 (双GPU)"
    print_message $CYAN "   - 显存<8GB: 推荐选择3 (三GPU)或更多GPU分布式训练"
    echo ""
    
    while true; do
        print_message $BLUE "请选择您需要的GPU配置 (输入对应数字):"
        read -p "输入选项编号 [1-单GPU/2-双GPU/3-三GPU/4-自定义/5-查看信息]: " gpu_choice
        
        case $gpu_choice in
            1)
                SELECTED_GPU_CONFIG="0"
                print_message $GREEN "✅ 选择: 单GPU模式 (GPU 0)"
                break
                ;;
            2)
                SELECTED_GPU_CONFIG="0,1"
                print_message $GREEN "✅ 选择: 双GPU模式 (GPU 0,1)"
                break
                ;;
            3)
                SELECTED_GPU_CONFIG="0,1,2"
                print_message $GREEN "✅ 选择: 三GPU模式 (GPU 0,1,2)"
                break
                ;;
            4)
                echo ""
                print_message $CYAN "自定义GPU配置示例:"
                echo "  0        - 使用GPU 0"
                echo "  0,1      - 使用GPU 0和1"
                echo "  1,2,3    - 使用GPU 1、2、3"
                echo "  2        - 使用GPU 2"
                echo ""
                read -p "请输入GPU编号 (用逗号分隔): " SELECTED_GPU_CONFIG
                
                if [ -z "$SELECTED_GPU_CONFIG" ]; then
                    print_message $YELLOW "⚠️  未输入GPU编号，默认使用GPU 0"
                    SELECTED_GPU_CONFIG="0"
                fi
                
                print_message $GREEN "✅ 选择: 自定义配置 (GPU $SELECTED_GPU_CONFIG)"
                break
                ;;
            5)
                echo ""
                print_message $CYAN "完整GPU列表:"
                if command -v nvidia-smi &> /dev/null; then
                    nvidia-smi --query-gpu=index,name,memory.total,memory.used --format=csv,noheader,nounits | while read line; do
                        if [ -n "$line" ]; then
                            IFS=',' read -r gpu_id name total_mem used_mem <<< "$line"
                            usage_percent=$(( used_mem * 100 / total_mem ))
                            print_message $YELLOW "  GPU $gpu_id: $name (${used_mem}MB/${total_mem}MB, ${usage_percent}%)"
                        fi
                    done
                else
                    print_message $YELLOW "  ⚠️  nvidia-smi不可用"
                fi
                echo ""
                ;;
            *)
                print_message $RED "❌ 无效选择，请输入1-5"
                ;;
        esac
    done
}

# 训练方式选择界面
select_training_stage() {
    print_message $BLUE "🎯 训练方式选择"
    print_message $CYAN "请选择训练模式:"
    print_message $YELLOW "┌──────────────────────────────────────────────────────────────┐"
    print_message $YELLOW "│  1. 完整两阶段训练 (强烈推荐)                        │"
    print_message $YELLOW "│     阶段1: 训练视觉-语言对齐层                      │"
    print_message $YELLOW "│     阶段2: 全模型联合训练                           │"
    print_message $YELLOW "│     时间: 约2-4小时 (根据GPU数量)                   │"
    print_message $YELLOW "│     效果: 最佳性能                                   │"
    print_message $YELLOW "│                                                              │"
    print_message $YELLOW "│  2. 仅阶段1训练 (Aligner层训练)                        │"
    print_message $YELLOW "│     适合: 初次训练、测试、内存不足                    │"
    print_message $YELLOW "│     时间: 约30-60分钟                               │"
    print_message $YELLOW "│     效果: 基础性能                                   │"
    print_message $YELLOW "│                                                              │"
    print_message $YELLOW "│  3. 仅阶段2训练 (需要阶段1结果)                       │"
    print_message $YELLOW "│     前置: 必须先完成阶段1训练                        │"
    print_message $YELLOW "│     时间: 约1-2小时                                 │"
    print_message $YELLOW "│     效果: 在阶段1基础上进一步优化                    │"
    print_message $YELLOW "└──────────────────────────────────────────────────────────────┘"
    
    while true; do
        print_message $BLUE "请选择您需要的训练模式 (输入对应数字):"
        read -p "输入选项编号 [1-完整训练/2-仅阶段1/3-仅阶段2]: " stage_choice
        
        case $stage_choice in
            1)
                SELECTED_STAGE_CONFIG="all"
                print_message $GREEN "✅ 选择: 完整两阶段训练"
                break
                ;;
            2)
                SELECTED_STAGE_CONFIG="1"
                print_message $GREEN "✅ 选择: 仅阶段1训练"
                break
                ;;
            3)
                # 检查阶段1结果是否存在
                if [ -d "models/universal_document/stage1" ] && [ -n "$(find models/universal_document/stage1 -name 'checkpoint-*' -type d 2>/dev/null)" ]; then
                    SELECTED_STAGE_CONFIG="2"
                    print_message $GREEN "✅ 选择: 仅阶段2训练"
                    break
                else
                    print_message $YELLOW "⚠️  未找到阶段1训练结果，建议先运行阶段1训练"
                    read -p "是否仍要继续? [y/N]: " confirm
                    if [[ $confirm =~ ^[Yy]$ ]]; then
                        SELECTED_STAGE_CONFIG="2"
                        print_message $GREEN "✅ 选择: 仅阶段2训练 (强制)"
                        break
                    fi
                fi
                ;;
            *)
                print_message $RED "❌ 无效选择，请输入1-3"
                ;;
        esac
    done
}

# 交互式配置界面
interactive_config() {
    print_message $PURPLE "🎛️  交互式训练配置"
    echo "====================================="
    
    # GPU配置选择 - 使用全局变量
    select_gpu_config
    local gpu_config="$SELECTED_GPU_CONFIG"
    echo ""
    
    # 训练方式选择 - 使用全局变量
    select_training_stage  
    local stage_config="$SELECTED_STAGE_CONFIG"
    echo ""
    
    # 确认配置
    print_message $CYAN "📋 配置确认:"
    echo "  🖥️  GPU配置: $gpu_config"
    echo "  🎯 训练方式: $stage_config"
    echo ""
    
    read -p "确认开始训练? [Y/n]: " confirm
    if [[ ! $confirm =~ ^[Nn]$ ]]; then
        print_message $GREEN "🚀 开始训练..."
        echo ""
        # 设置全局变量供主函数使用
        FINAL_GPU_CONFIG="$gpu_config"
        FINAL_STAGE_CONFIG="$stage_config"
        return 0
    else
        print_message $YELLOW "❌ 用户取消训练"
        exit 0
    fi
}

# 主函数
main() {
    # 默认参数
    local stage="interactive"
    local gpu_ids="interactive"
    local monitor_only=false
    local tensorboard_only=false
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            --stage)
                stage="$2"
                shift 2
                ;;
            --gpu)
                gpu_ids="$2"
                shift 2
                ;;
            --monitor-only)
                monitor_only=true
                shift
                ;;
            --tensorboard)
                tensorboard_only=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                print_message $RED "未知选项: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    # 打印标题
    print_title
    
    # 仅启动TensorBoard
    if [ "$tensorboard_only" = true ]; then
        start_tensorboard
        exit 0
    fi
    
    # 检查环境
    check_environment
    
    # 显示GPU信息
    show_gpu_info
    
    # 启动TensorBoard
    start_tensorboard
    
    # 仅监控模式
    if [ "$monitor_only" = true ]; then
        print_message $CYAN "📊 监控模式：等待现有训练进程..."
        # 这里可以添加监控现有进程的逻辑
        exit 0
    fi
    
    # 交互式配置模式
    if [ "$stage" = "interactive" ] || [ "$gpu_ids" = "interactive" ]; then
        # 获取用户配置
        interactive_config
        if [ $? -eq 0 ]; then
            # 使用全局变量获取配置结果
            gpu_ids="$FINAL_GPU_CONFIG"
            stage="$FINAL_STAGE_CONFIG"
            print_message $BLUE "📋 使用配置: GPU=$gpu_ids, 阶段=$stage"
        else
            print_message $RED "❌ 配置失败，退出"
            exit 1
        fi
    fi
    
    # 执行训练
    case $stage in
        1)
            run_training "1" "$gpu_ids"
            ;;
        2)
            run_training "2" "$gpu_ids"
            ;;
        all)
            print_message $PURPLE "🎯 开始完整两阶段训练流程"
            if run_training "1" "$gpu_ids"; then
                print_message $GREEN "✅ 阶段1训练成功完成！"
                print_message $BLUE "===================================================="
                print_message $BLUE "🚀 准备开始阶段2训练"
                print_message $BLUE "===================================================="
                
                # 清理GPU缓存和等待资源释放
                print_message $YELLOW "🧹 清理GPU缓存和资源..."
                python -c "import torch; torch.cuda.empty_cache(); print('✅ GPU缓存已清理')" 2>/dev/null || true
                
                print_message $BLUE "⏳ 等待10秒后开始阶段2训练，请稍候..."
                for i in {10..1}; do
                    print_message $CYAN "   倒计时: ${i} 秒"
                    sleep 1
                done
                print_message $GREEN "🚀 现在开始阶段2训练！"
                
                if run_training "2" "$gpu_ids"; then
                    print_message $GREEN "🎉 完整两阶段训练成功完成！"
                else
                    print_message $RED "❌ 阶段2训练失败"
                    exit 1
                fi
            else
                print_message $RED "❌ 阶段1失败，终止训练流程"
                exit 1
            fi
            ;;
        *)
            print_message $RED "无效的阶段参数: $stage"
            show_help
            exit 1
            ;;
    esac
    
    print_message $GREEN "🎉 训练流程完成！"
    
    # 提供后续操作建议
    if [ -d "models/universal_document" ]; then
        echo ""
        print_message $CYAN "🎉 训练完成！后续操作:"
        echo "  📊 查看训练日志: tensorboard --logdir logs/universal_document"
        echo "  🔮 测试推理: python scripts/vllm_inference.py --gpu_id 3 --model models/universal_document/stage2 --image data/test_images/contract_01.jpeg"
        echo "  🌐 启动API: python scripts/deploy_service.py --gpu_id 3 --model models/universal_document/stage2"
        echo "  📈 评估模型: python scripts/evaluation.py --gpu_id 3 --predictions data/test_predictions.json --ground_truth data/test_ground_truth.json"
    fi
}

# 信号处理
trap 'print_message $YELLOW "⏹️  训练被中断"; exit 1' INT TERM

# 运行主函数
main "$@"